package com.zxy.leetcode.lcof;

import java.util.Arrays;

/**
 * https://leetcode-cn.com/problems/zui-xiao-de-kge-shu-lcof/
 *
 * 最小的k个数
 *
 * 输入整数数组 arr ，找出其中最小的 k 个数。
 * 例如，输入4、5、1、6、2、7、3、8这8个数字，则最小的4个数字是1、2、3、4。
 *
 * 限制：
 * 0 <= k <= arr.length <= 10000
 * 0 <= arr[i] <= 10000
 *
 * 标签：分治、快排
 *
 * 与 Test00215（数组中的第K个最大元素）解法类似，用快排的思想，找出第K小的位置即可
 */
public class Test40 {

    public static void main(String[] args) {
        Test40 test = new Test40();

        int[] nums1 = {3,2,1};
        System.out.println(Arrays.toString(test.getLeastNumbers(nums1, 2)));
        System.out.println(Arrays.toString(test.getLeastNumbers(nums1, 1)));

        int[] nums2 = {0,1,2,1};
        System.out.println(Arrays.toString(test.getLeastNumbers(nums2, 1)));

    }

    public int[] getLeastNumbers(int[] arr, int k) {
        if (arr.length == k) {
            return arr;
        }
        if (k == 0) {
            return new int[0];
        }

        partition(arr, 0, arr.length-1, k);

        return Arrays.copyOf(arr, k);
    }

    private void partition(int[] nums, int left, int right,int k) {
        // 防止nums刚好是逆序的
        swap(nums, left, (left + right)/2);

        int pivot = nums[left]; // 中心轴
        int i = left;
        int j = right;
        while (i < j) {
            while (i < j && nums[j] >= pivot) j--;
            swap(nums, i, j);

            while (i < j && nums[i] <= pivot) i++;
            swap(nums, i, j);
        }

        int pivotIndex = i;
        if (pivotIndex + 1 == k) {
            return;
        }
        if (pivotIndex + 1 > k) {
            partition(nums, left, pivotIndex-1, k);
        }else {
            partition(nums, pivotIndex+1, right, k);
        }
    }

    private void swap(int[] nums, int a, int b) {
        if (a == b) {
            return;
        }
        int tmp = nums[a];
        nums[a] = nums[b];
        nums[b] = tmp;
    }

}
